import numpy as np
from sort_API import parse_page_data
from config import N_CLUSTERS

from sklearn.cluster import KMeans


def kmeans_up_down_split(page_data):
    """
    使用K-means进行上下分栏，处理样本数小于等于聚类数的情况
    """
    coords, _, _ = parse_page_data(page_data)

    # 如果没有坐标数据，直接返回空列表
    if not coords:
        print("K-means分栏: 无坐标数据，返回空列表")
        return []

    # 提取y坐标
    y_centers = np.array([item['y_center'] for item in coords]).reshape(-1, 1)
    n_samples = len(y_centers)

    # 确定实际使用的聚类数（不超过样本数）
    actual_clusters = min(N_CLUSTERS, n_samples)

    # 当样本数小于等于1时，直接返回按y_start排序的原始数据作为单栏
    if n_samples <= 1:
        coords.sort(key=lambda x: x['y_start'])
        print(f"K-means分栏: 样本数较少({n_samples}个)，不进行聚类，返回1栏")
        return [coords]

    # 使用K-means聚类
    try:
        kmeans = KMeans(n_clusters=actual_clusters, random_state=42, n_init=10)
        labels = kmeans.fit_predict(y_centers)
    except Exception as e:
        print(f"K-means聚类出错: {str(e)}，返回原始排序结果")
        coords.sort(key=lambda x: x['y_start'])
        return [coords]

    # 分组
    columns = {}
    for i, (item, label) in enumerate(zip(coords, labels)):
        if label not in columns:
            columns[label] = []
        columns[label].append(item)

    # 按y_start排序每个栏内的元素（从上到下）
    sorted_columns = []
    for label, column_data in columns.items():
        column_data.sort(key=lambda x: x['y_start'])
        sorted_columns.append(column_data)

    # 按平均y坐标排序栏（从上到下）
    sorted_columns.sort(key=lambda col: np.mean([item['y_center'] for item in col]))

    print(f"K-means分栏: 共{len(sorted_columns)}栏")
    for i, column in enumerate(sorted_columns):
        avg_y = np.mean([item['y_center'] for item in column])
        print(f"  栏 {i + 1}: {len(column)}字, 平均Y位置: {avg_y:.1f}")

    return sorted_columns
